#pragma once
#include<iostream>
#include<cstring>
#include<opencv2/core.hpp>
#include<vector>
#include<opencv2/opencv.hpp>
#include<opencv2/highgui.hpp>
#include<opencv2/imgcodecs.hpp>
using namespace std;
using namespace cv;
float w=132,h=62;//单位是mm，测出的装甲版大概宽度高度
vector<Point3f> p3d ={Point3d(w/2,h/2,0),Point3d(-w/2,h/2,0),Point3d(-w/2,-h/2,0),Point3d(w/2,-h/2,0)};

class solver{
public:
    Mat cam_matrix = (Mat_<double>(3, 3) << 2673.8, 0, 718.2889,
            0, 2673.2, 560.3110,
            0, 0, 1);                // 内参矩阵
    Mat distortion_coeff = (Mat_<double>(5, 1) << -0.2097, -0.4538, 0, 0, 0);          // 畸变系数
    Mat rvec,tvec;
    vector<RotatedRect>recs;
    Point2f rec1[4],rec2[4],rec3[4],rec[4];
    vector<Point2d> p2d;
    inline double len(Mat vec){
        double tot=0;
        for(int i=0;i<3;i++)
            tot += vec.at<double>(0,i) * vec.at<double>(0,i);
        return sqrt(tot);
    }
    inline void draw_poly(Point2f rec[4],Mat &img){
        for(int i=0;i<4;i++)
            line(img,rec[i],rec[(i+1)%4],Scalar(0,255,0),3,LINE_AA);
    }
    static inline bool cmp_ro(RotatedRect &a,RotatedRect &b){return a.center.y<b.center.y;}
    static inline bool cmp(Point2f &a,Point2f &b){return a.y<b.y;}
    inline void get2D(Mat &frame){
            recs[0].points(rec1),recs[1].points(rec2);
            sort(rec1,rec1+4,cmp),sort(rec2,rec2+4,cmp);
            Point2f up1=(rec1[0]+rec1[1])/2,dw1=(rec1[2]+rec1[3])/2,up2=(rec2[0]+rec2[1])/2,dw2=(rec2[2]+rec2[3])/2;
            if(up1.x>up2.x)swap(up1,up2),swap(dw1,dw2);
            Point2f rec[4]={up1,up2,dw2,dw1};
            p2d.clear(),p2d.push_back(up1),p2d.push_back(up2),p2d.push_back(dw2),p2d.push_back(dw1);
            draw_poly(rec,frame);

            {solvePnP(p3d, p2d, cam_matrix, distortion_coeff,rvec,tvec, 0,cv::SOLVEPNP_AP3P);}
            putText(frame,"d= "+to_string(len(tvec)/1000)+"m",dw2, 2,1,Scalar(0,0,255),1,LINE_AA);
    }
    inline void get3D(Mat &frame){
        sort(recs.begin(),recs.end(),cmp_ro);
        if(abs(recs[0].center.x-recs[1].center.x)>abs(recs[1].center.x-recs[2].center.x))swap(recs[0],recs[2]);
        get2D(frame);
    }
    inline void get4D(Mat &frame){
        sort(recs.begin(),recs.end(),cmp_ro);
        get2D(frame);
        swap(recs[0],recs[2]),swap(recs[1],recs[3]);
        get2D(frame);
    }
    void detect(cv::VideoCapture video,int step) {

        cv::Mat frame;
        cv::Mat struct1 = getStructuringElement(0, cv::Size(3, 3));
        vector<vector<cv::Point>> coun_re;

        int fs = video.get(CAP_PROP_FRAME_COUNT);
        int start = clock();
        double total_time = 0;
        while (video.read(frame)) {
            int start = clock();
            recs.clear();
            Mat gray;
            cvtColor(frame,gray,COLOR_BGR2GRAY);
            threshold(gray,gray,110,255,THRESH_BINARY);
            GaussianBlur(gray,gray,Size(3,3),1);
            gray(Rect(0,0,1440,800))=0;
            findContours(gray,coun_re,RETR_TREE,CHAIN_APPROX_NONE);
        //  dilate(gray,gray,struct1,Point(-1,-1),1);
            for(int i=0;i<coun_re.size();i++) {
                double area= contourArea(coun_re[i]);
                if(area<140||area>2000)continue;
                RotatedRect t= minAreaRect(coun_re[i]);
                t.points(rec1);
                if(max(t.size.width,t.size.height)<50)continue;
                memcpy(rec2,rec1,sizeof rec2);
                sort(rec1,rec1+4,cmp);
                if(abs(rec1[0].x-rec1[1].x)>50)continue;
                recs.push_back(t);
            }
            if(recs.size()==2)get2D(frame);
            else if(recs.size()==3)get3D(frame);
            else get4D(frame);
            int end=clock();
            total_time+=((double)end-start)/CLOCKS_PER_SEC;
            imshow("s",frame);
            waitKey(step);
        }
        int end = clock();
        double total = (double) (end - start) / CLOCKS_PER_SEC - (double) fs * step / 1000;

        printf("一共%d张图片，花费%lf秒，平均每张图片%lf秒 实际算法执行时间每张图%lf秒\n", fs, total, (double) total / fs,
            total_time / fs);

    }
};
